import os
from pathlib import Path
from PyPDF2 import PdfReader
from PIL import Image
import io
import sys
# 获取当前文件的目录
current_dir = os.path.dirname(os.path.abspath(__file__))
# 添加上级目录到系统路径
sys.path.append(os.path.dirname(current_dir))
from image_to_text import image_to_text

def convert_pdf_to_images(pdf_path, output_dir=None, dpi=300, ocr_method='default'):
    try:
        # 确保PDF文件存在
        pdf_path = Path(pdf_path)
        if not pdf_path.exists():
            print(f"错误：PDF文件 '{pdf_path}' 不存在")
            return False
        
        # 设置输出目录
        if output_dir is None:
            output_dir = pdf_path.parent / pdf_path.stem
        else:
            output_dir = Path(output_dir)
        
        # 创建输出目录
        output_dir.mkdir(exist_ok=True, parents=True)
        
        print(f"开始处理PDF文件：{pdf_path}")
        print(f"输出目录：{output_dir}")
        
        # 读取PDF文件
        reader = PdfReader(pdf_path)
        total_pages = len(reader.pages)
        print(f"PDF共有 {total_pages} 页")
        
        # 处理每一页
        for i in range(total_pages):
            page = reader.pages[i]
            current_page = i + 1  # 添加当前页码计数（从1开始）
            
            # 提取文本内容
            text = page.extract_text()
            if text.strip():  # 如果有文本内容
                # 在文件名前添加页码保证顺序
                text_file = output_dir / f"{current_page:03d}_text_{i:03d}.txt"
                text_file.write_text(text, encoding='utf-8')
                print(f"已保存第 {current_page} 页文本：{text_file}")
            
            # 提取页面中的图片（添加图片序号）
            for j, image_file_object in enumerate(page.images, 1):
                # 保存图片
                image = Image.open(io.BytesIO(image_file_object.data))
                # 文件名格式：页码_图片序号.扩展名
                output_path = output_dir / f"{current_page:03d}_image_{j:02d}{Path(image_file_object.name).suffix}"
                image.save(output_path)
                print(f"已保存第 {current_page} 页图片：{output_path}")
                
                # 从图片中提取文字（保持相同序号）
                if ocr_method == 'ai':
                    from image_to_text_ai_cli import image_to_text_ai
                    extracted_text = image_to_text_ai(str(output_path))
                else:
                    extracted_text = image_to_text(output_path)
                if extracted_text:
                    # 保持文件名前缀与图片一致
                    text_file = output_dir / f"{current_page:03d}_image_{j:02d}_text.txt"
                    text_file.write_text(extracted_text, encoding='utf-8')
                    print(f"已保存第 {current_page} 页图片文字：{text_file}")
        
        print("\n转换完成！")
        print(f"所有内容已保存到：{output_dir}")
        return True
        
    except Exception as e:
        print(f"转换过程中出错：{str(e)}")
        import traceback
        print(f"错误详情：\n{traceback.format_exc()}")
        return False

def process_directory(dir_path, dpi=300, ocr_method='default'):
    """处理目录下的所有PDF文件
    
    Args:
        dir_path: 目录路径
        dpi: 输出图片的DPI
    
    Returns:
        tuple: (成功数量, 失败数量)
    """
    dir_path = Path(dir_path)
    if not dir_path.is_dir():
        print(f"错误：'{dir_path}' 不是有效目录")
        return (0, 0)
    
    # 查找所有PDF文件
    pdf_files = list(dir_path.glob('**/*.pdf'))
    if not pdf_files:
        print("未找到任何PDF文件")
        return (0, 0)
    
    print(f"找到 {len(pdf_files)} 个PDF文件")
    success_count = 0
    failed_count = 0
    
    # 处理每个PDF文件
    for i, pdf_file in enumerate(pdf_files, 1):
        print(f"\n正在处理 ({i}/{len(pdf_files)}): {pdf_file.name}")
        output_dir = pdf_file.parent / pdf_file.stem
        
        if convert_pdf_to_images(pdf_file, output_dir, dpi, ocr_method):
            success_count += 1
        else:
            failed_count += 1
    
    print(f"\n处理完成！成功：{success_count}，失败：{failed_count}")
    return (success_count, failed_count)

def main():
    print("PDF图片提取工具")
    print("-" * 20)
    
    # 获取用户输入
    while True:
        path = input("请输入PDF文件或目录路径：").strip()
        if not path:
            print("请输入有效路径")
            continue
            
        path = Path(path)
        if path.is_file() and path.suffix.lower() == '.pdf':
            break
        elif path.is_dir():
            break
        else:
            print("请输入有效的PDF文件或目录路径")
    
    # 获取DPI设置（可选）
    while True:
        dpi_input = input("请输入图片DPI（可选，直接回车使用默认值300，建议范围300-1200）：").strip()
        if not dpi_input:
            dpi = 300
            break
        try:
            dpi = int(dpi_input)
            if dpi > 0:
                if dpi < 300:
                    print("警告：DPI值过低可能影响图像质量，建议使用300以上的值")
                elif dpi > 1200:
                    print("警告：DPI值过高可能导致处理时间较长和文件体积较大")
                break
            print("DPI必须大于0")
        except ValueError:
            print("请输入有效的数字")
    
        # 获取OCR方式选择
    while True:
        ocr_method = input("请选择OCR方式（default/ai，默认default）：").strip().lower()
        if not ocr_method:
            ocr_method = 'default'
        if ocr_method in ['default', 'ai']:
            break
        print("请输入有效的OCR方式（default 或 ai）")

    print("\n开始提取，请稍候...")
    
    # 根据输入路径类型执行相应的转换
    if path.is_file():
        convert_pdf_to_images(path, None, dpi)
    else:
        process_directory(path, dpi,ocr_method)

if __name__ == "__main__":
    main()